import torch  
import torch.nn as nn  
import torchvision.transforms as transforms  
from torchvision import models  
from PIL import Image  
import matplotlib.pyplot as plt  

# 功能函数：加载图像并进行预处理  
def load_image(image_path, max_size=400, shape=None):  
    """  
    加载图像并转换为张量。  
    
    参数：  
    - image_path: 图像文件路径  
    - max_size: 最大边长（宽或高），以保持图像比例  
    - shape: 指定输出图像的形状（高度和宽度）  
    
    返回：  
    处理后的图像张量（形状为 [1, C, H, W]）  
    """  
    # 打开图像并转换为 RGB 格式  
    image = Image.open(image_path).convert("RGB")  
    
    # 如果指定最大尺寸，调整图像大小  
    if max_size is not None:  
        size = max(max(image.size), max_size)  # 确保最大边不超过 max_size  
        image = transforms.Resize(size)(image)  

    # 如果给定形状，则调整图像到特定形状  
    if shape is not None:  
        image = transforms.Resize(shape)(image)  

    # 将图像转换为张量，并增加批量维度  
    transform = transforms.ToTensor()  
    image = transform(image).unsqueeze(0)  
    return image  

# 功能函数：显示图像  
def imshow(tensor, title=None):  
    """  
    显示图像张量。（由 torch.Tensor 转换为 PIL Image）  
    
    参数：  
    - tensor: 输入图像的张量  
    - title: 图像标题  
    """  
    image = tensor.cpu().clone()  # 克隆张量以避免修改原始内容  
    image = image.squeeze(0)  # 去掉批量维度 [1, C, H, W] --> [C, H, W]  
    image = transforms.ToPILImage()(image)  # 将张量转换为 PIL 图像  
    plt.imshow(image)  
    if title is not None:  
        plt.title(title)  # 可选标题  
    plt.axis('off')  # 不显示坐标轴  
    plt.show()  

# 加载内容图像和风格图像  
content_img = load_image("path/to/content_image.jpg", max_size=400)  # 内容图像  
style_img = load_image("path/to/style_image.jpg", shape=content_img.shape[-2:])  # 风格图像，形状与内容图像相同  

# 使用 VGG19 模型提取特征  
vgg = models.vgg19(pretrained=True).features  # 加载预训练的 VGG19 模型  
for param in vgg.parameters():  
    param.requires_grad = False  # 冻结参数，避免在训练过程中更新  

# 检查使用的设备（GPU 或 CPU）  
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
vgg = vgg.to(device)  # 将 VGG 模型移动到计算设备  

# 将图像传入 VGG 网络以获取特征  
content_img = content_img.to(device)  # 将内容图像移动到设备  
style_img = style_img.to(device)  # 将风格图像移动到设备  

def get_features(model, layers, x):  
    """  
    从 VGG 模型中提取特征。  
    
    参数：  
    - model: 预训练模型  
    - layers: 要提取特征的层的名称列表  
    - x: 输入图像张量  
    
    返回：  
    特征字典，键为层名，值为相应的特征图  
    """  
    features = {}  
    # 遍历模型中的每一层  
    for name, layer in model._modules.items():  
        x = layer(x)  # 将输入图像通过当前层  
        if name in layers:  
            features[name] = x  # 如果当前层在指定层中，保存特征图  
    return features  

# 定义需要提取特征的层  
content_layers = ['21']  # 选择 VGG19 的 conv4_2 作为内容特征  
style_layers = ['0', '5', '10', '19', '28']  # 选择 VGG19 的多个卷积层作为风格特征  
layers = content_layers + style_layers  # 合并内容层与风格层  

# 获取内容和风格特征  
content_features = get_features(vgg, content_layers, content_img)  # 从内容图像中提取特征  
style_features = get_features(vgg, style_layers, style_img)  # 从风格图像中提取特征  

# 计算风格损失时使用的 Gram 矩阵  
def gram_matrix(tensor):  
    """  
    计算 Gram 矩阵，通常用于计算风格损失。  
    
    参数：  
    - tensor: 需要计算的特征图，通常为 [C, H, W]  
    
    返回：  
    Gram 矩阵，归一化后的风格特征  
    """  
    b, c, h, w = tensor.size()  # 获取张量的形状  
    features = tensor.view(b, c, h * w)  # 将特征图展开为 [B, C, H*W]  
    gram = torch.bmm(features, features.transpose(1, 2))  # 计算 Gram 矩阵  
    return gram.div(b * c * h * w)  # 归一化  

# 计算风格图像的 Gram 矩阵  
style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_layers}  

# 创建目标图像，使用内容图像作为初始化  
target = content_img.clone().requires_grad_(True)  # 允许对目标图像进行梯度计算  

# 设置损失权重  
content_weight = 1e5  # 内容损失权重  
style_weight = 1e10   # 风格损失权重  

# 定义优化器以更新目标图像  
optimizer = torch.optim.Adam([target], lr=0.003)  # 使用 Adam 优化器  

# 风格迁移过程  
num_iterations = 500  # 设置迭代次数  
for i in range(num_iterations):  
    # 获取当前目标图像的特征  
    target_features = get_features(vgg, layers, target)  

    # 计算内容损失  
    content_loss = nn.MSELoss()(target_features[content_layers[0]], content_features[content_layers[0]])  

    # 计算风格损失  
    style_loss = 0  
    for layer in style_layers:  
        target_gram = gram_matrix(target_features[layer])  # 当前目标图像在该层的 Gram 矩阵  
        style_loss += nn.MSELoss()(target_gram, style_grams[layer])  # 计算风格损失  

    # 总损失 = 内容损失 + 风格损失  
    total_loss = content_weight * content_loss + style_weight * style_loss  

    optimizer.zero_grad()  # 清空之前的梯度  
    total_loss.backward()  # 反向传播计算梯度  
    optimizer.step()  # 更新目标图像  

    # 每 50 次输出一次损失  
    if i % 50 == 0:  
        print(f"Iteration {i}, Total Loss: {total_loss.item()}")  

# 显示结果图像  
imshow(target, title='Generated Image')